package limiter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

/**
 * @author dushougudu
 */
public class LeakybucketLimiter {
    /**
     * 每秒处理数（出水率）
     */
    private volatile long rate;

    /**
     * 当前剩余水量
     */
    private volatile long currentWater;

    /**
     * 最后刷新时间
     */
    private volatile long refreshTime;

    /**
     * 桶容量
     */
    private volatile long capacity;


    public LeakybucketLimiter(long capacity, long rate) {
        this.capacity = capacity;
        this.rate = rate;
    }

    /**
     * 漏桶算法
     *
     * @return
     */
    public synchronized boolean leakybucketLimitTryAcquire() {
        //获取系统当前时间
        long currentTime = System.currentTimeMillis();
        //流出的水量 =(当前时间-上次刷新时间)* 出水率
        long outWater = (currentTime - refreshTime) / 1000 * rate;
        // 当前水量 = 之前的桶内水量-流出的水量
        currentWater = Math.max(0, currentWater - outWater);
        // 刷新时间
        refreshTime = currentTime;

        // 当前剩余水量还是小于桶的容量，则请求放行
        if (currentWater < capacity) {
            currentWater++;
            return true;
        }
        // 当前剩余水量大于等于桶的容量，限流
        return false;
    }


    public static void main(String[] args) throws ExecutionException, InterruptedException {
        LeakybucketLimiter leakybucketLimiter = new LeakybucketLimiter(5, 3);

        CyclicBarrier barrier = new CyclicBarrier(5);
        List<CompletableFuture<Boolean>> list = new ArrayList<>();
        ThreadPoolExecutor executor = new ThreadPoolExecutor(10, 10,
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>());
        for (int i = 0; i < 10; i++) {
            if (i == 5) {
                Thread.sleep(1000);
            }
            CompletableFuture<Boolean> future = CompletableFuture.supplyAsync(() -> {
                try {
                    barrier.await();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } catch (BrokenBarrierException e) {
                    e.printStackTrace();
                }
                boolean flag = leakybucketLimiter.leakybucketLimitTryAcquire();
                return flag;
            }, executor);
            list.add(future);
        }

        for (CompletableFuture<Boolean> future : list) {
            System.out.println(future.get());
        }

        executor.shutdown();
    }

}
